import sys
from pickle import load

import numpy as np
import pandas as pd
from onnxruntime import InferenceSession
from sklearn.metrics import mean_absolute_error

from constant import DatasetTestPath,LabelScalerPath,DatasetTrainPath,\
    CHECK_POINT,TESTB_PATH,ANS_PATH,OUTPUT_PATH

ONNX_MODEL = f'{CHECK_POINT}DLPriceModel_6_1_5_1-4.onnx'
if len(sys.argv) == 2:
    ONNX_MODEL = sys.argv[1]
DL_MODEL_NAME = ONNX_MODEL.split('.')[-2].split('\\')[-1]
print('model name', DL_MODEL_NAME)
print('model location', ONNX_MODEL)

datasetTest = np.load(DatasetTestPath).astype(np.float32)
datasetTrain = np.load(DatasetTrainPath).astype(np.float32)
print('datasetTest.shape',datasetTest.shape)
print('datasetTrain.shape',datasetTrain.shape)
with open(LabelScalerPath,'br') as f:
    label_scaler = load(f)
print('data loaded.')

model = InferenceSession(ONNX_MODEL)
print('model loaded.')

out = model.run(
        # [model.get_outputs()[0].name],
        None,
        {model.get_inputs()[0].name:datasetTrain})[0]
predict = pd.read_csv(f'{OUTPUT_PATH}predict.csv')
predict['predict_raw'] = out
predict['predict'] = label_scaler.inverse_transform(out)
score = mean_absolute_error(predict['label'].values,predict['predict'].values)
predict.to_csv(f'{OUTPUT_PATH}predict.csv', index=False)
print(f'predict trainset finished with score {score}!')

# 预测
ans = model.run(
        [model.get_outputs()[0].name],
        {model.get_inputs()[0].name:datasetTest})[0]
testID = pd.read_csv(TESTB_PATH, sep=' ')['SaleID']
re=pd.DataFrame()
re['SaleID'] = testID
re['price'] = label_scaler.inverse_transform(ans)
re['price'] = re['price'].map(lambda x: 0 if x<0 else x)
re.to_csv(f'{ANS_PATH}ans_{DL_MODEL_NAME}_{score}.csv', index=False)
print('predict testset finished!')